//
//  ipv4_defragmentor.cpp
//  testpcap
//
//  Created by bugforever on 1/15/15.
//  Copyright (c) 2015 bugforever. All rights reserved.
//

#include "ipv4_defragmentor.h"
#include "string.h"
#include "stdio.h"
#include "../utility/os_api.h"

#define IP_FRAGEMENT_TIMEOUT ( 30 * 1000 )

class ipv4_fragment_cmp
{
public:
    bool operator () (const ipv4_fragment * ptr_f1 ,const ipv4_fragment * ptr_f2)
    {
        return ptr_f1->fragment_offset < ptr_f2->fragment_offset;
    }
};

ipv4_defragmentor::ipv4_defragmentor()
{
    ;
}

ipv4_defragmentor::~ipv4_defragmentor()
{
    destroy_fragments_map();
}

int ipv4_defragmentor::defragment( ipv4_data * ptr_input , network_ipv4_data * ptr_output )
{
    if( !ptr_input || !ptr_output )
        return -1;
    
    if( !ptr_input->will_fragment )
    {
        copy_ipv4data(ptr_input, ptr_output);
        return 0;
    }
    
    if( ptr_input->is_last_fragment && ptr_input->fragment_offset == 0 )
    {
        copy_ipv4data(ptr_input, ptr_output);
        return 0;
    }
    
    printf( "The IP data is fragmented!!!\n" );
    
    ipv4_fragment_list * ptr_list = add_fragment( ptr_input );
    
    clean_timeout_fragmentlist();
    
    bool is_complete = is_fragmentlist_complete( ptr_list );

    if( is_complete )
    {
        defragment_fragmentlist( ptr_list , ptr_output );
        
        destroy_fragmentlist( ptr_list );
        m_fragments_map.erase( ptr_input->identification );
        
        return 0;
    }
    
    return -2;
}

void ipv4_defragmentor::copy_ipv4data( ipv4_data * ptr_input , network_ipv4_data * ptr_output )
{
    memcpy( ptr_output->saddr , ptr_input->saddr , 4 );
    memcpy( ptr_output->daddr , ptr_input->daddr , 4 );
    ptr_output->ptr_data = ptr_input->ptr_data;
}

ipv4_fragment * ipv4_defragmentor::create_fragment( ipv4_data * ptr_data )
{
    ipv4_fragment * ptr_fragment = new ipv4_fragment;
    memcpy( ptr_fragment->saddr , ptr_data->saddr , 4 );
    memcpy(ptr_fragment->daddr, ptr_data->daddr, 4 );
    ptr_fragment->is_last_fragment = ptr_data->is_last_fragment;
    ptr_fragment->fragment_offset = ptr_data->fragment_offset;
    ptr_fragment->ptr_data = ptr_data->ptr_data;
    
    return ptr_fragment;
}

ipv4_fragment_list * ipv4_defragmentor::add_fragment( ipv4_data * ptr_data )
{
    ipv4_fragment * ptr_fragment = create_fragment( ptr_data );
    
    ipv4_fragment_list * ptr_list = 0;
    
    if( m_fragments_map.find( ptr_data->identification ) == m_fragments_map.end() )
    {
        m_fragments_map[ ptr_data->identification ] = ptr_list = new ipv4_fragment_list;
        ptr_list->latest_fragment_time = 0;
        ptr_list->ptr_list = new std::list< ipv4_fragment * >;
    }
    else
        ptr_list = m_fragments_map[ ptr_data->identification ];
    
    ptr_list->latest_fragment_time = get_time();
    ptr_list->ptr_list->push_back(ptr_fragment);
    
    ptr_list->ptr_list->sort( ipv4_fragment_cmp() );
    
    return ptr_list;
}

bool ipv4_defragmentor::is_fragmentlist_complete( ipv4_fragment_list * ptr_list )
{
    std::list< ipv4_fragment * >::iterator iter = ptr_list->ptr_list->begin();
    std::list< ipv4_fragment * >::iterator itere = ptr_list->ptr_list->end();
    
    bool is_complete = false;
    int next_offset = 0;
    
    while( iter != itere )
    {
        if( next_offset != (*iter)->fragment_offset )
            break;
        
        if( (*iter)->is_last_fragment )
            is_complete = true;
        
        next_offset = next_offset + (*iter)->ptr_data->get_datasize_by_tag(NDT_IP_DATA);
        iter++;
    }
    
    return is_complete;
}

void ipv4_defragmentor::clean_timeout_fragmentlist()
{
    std::map< int , ipv4_fragment_list * >::iterator iter = m_fragments_map.begin();
    std::map< int , ipv4_fragment_list * >::iterator itere = m_fragments_map.end();
    
    while( iter != itere )
    {
        if( get_time() - iter->second->latest_fragment_time > IP_FRAGEMENT_TIMEOUT )
        {
            destroy_fragmentlist( iter->second );
            delete iter->second;
            iter = m_fragments_map.erase(iter);
        }
        else
            iter++;
    }
}

void ipv4_defragmentor::destroy_fragmentlist( ipv4_fragment_list * ptr_fragment_list )
{
    if( !ptr_fragment_list )
        return;
    
    std::list< ipv4_fragment * > * ptr_list = ptr_fragment_list->ptr_list;
    
    std::list< ipv4_fragment * >::iterator iter = ptr_list->begin();
    std::list< ipv4_fragment * >::iterator itere = ptr_list->end();
    
    while( iter != itere )
    {
        (*iter)->ptr_data->release();
        delete (*iter);
        
        iter = ptr_list->erase(iter);
    }
}

void ipv4_defragmentor::defragment_fragmentlist( ipv4_fragment_list * ptr_list , network_ipv4_data * ptr_output )
{
    std::list< ipv4_fragment * >::iterator iter = ptr_list->ptr_list->begin();
    std::list< ipv4_fragment * >::iterator itere = ptr_list->ptr_list->end();
    
    if( iter == itere )
        return;
    
    memcpy( ptr_output->saddr , (*iter)->saddr , 4);
    memcpy( ptr_output->daddr , (*iter)->daddr , 4);
    ptr_output->ptr_data = (*iter)->ptr_data->duplicate();
    
    iter++;
    
    while( iter != itere )
    {
        ptr_output->ptr_data->append_data((*iter)->ptr_data->get_data_by_tag(NDT_IP_DATA), (*iter)->ptr_data->get_datasize_by_tag(NDT_IP_DATA));
        
        iter++;
    }
}

void ipv4_defragmentor::destroy_fragments_map()
{
    std::map< int , ipv4_fragment_list * >::iterator iter = m_fragments_map.begin();
    std::map< int , ipv4_fragment_list * >::iterator itere = m_fragments_map.end();
    
    while( iter != itere )
    {
        destroy_fragmentlist( iter->second );
        delete iter->second;
        iter = m_fragments_map.erase(iter);
    }
}







